from pathlib import Path

from ipyniivue import download_dataset

BASE_API_URL = "https://niivue.com/demos/images"
DATA_FOLDER = Path("images")

# Download data for example
download_dataset(
    BASE_API_URL,
    DATA_FOLDER,
    files=[
        "BrainMesh_ICBM152.lh.mz3",
        "CIT168.mz3",
        "BrainMesh_ICBM152.lh.motor.mz3",
    ],
)

BrainMesh_ICBM152.lh.mz3 already exists.
Downloading CIT168.mz3...
BrainMesh_ICBM152.lh.motor.mz3 already exists.
Dataset downloaded successfully to images.

import ipywidgets as widgets
from IPython.display import display

from ipyniivue import Mesh, MeshLayer, NiiVue

# 1. Initialize NiiVue
nv = NiiVue()
nv.opts.show_3d_crosshair = True
nv.opts.back_color = (1, 1, 1, 1)
nv.opts.mesh_xray = 0.3
nv.opts.show_legend = False
nv.opts.is_colorbar = True

# 2. Define the initial meshes
mesh_hemisphere = Mesh(
    path=DATA_FOLDER / "BrainMesh_ICBM152.lh.mz3",
    rgba255=[255, 255, 255, 255],
    layers=[],
)

mesh_cit = Mesh(path=DATA_FOLDER / "CIT168.mz3", rgba255=[0, 0, 255, 255])

# Load meshes into the scene
nv.meshes = [mesh_hemisphere, mesh_cit]

# 3. Create UI Controls

# XRay Slider
xray_slider = widgets.FloatSlider(
    min=0, max=0.5, value=0.3, step=0.01, description="XRay"
)

# Hemisphere Visible Checkbox
vis_checkbox = widgets.Checkbox(value=True, description="Hemisphere Visible")

# Add Layer Buttons
clear_layers_btn = widgets.Button(description="Clear Layers")
add_layer_btn = widgets.Button(description="Add Layer")

# Layer Threshold Slider (cal_min)
layer_thresh_slider = widgets.FloatSlider(
    min=0.1, max=4.9, value=0.5, step=0.1, description="Layer Thresh"
)

# Layer Opacity Slider
layer_alpha_slider = widgets.FloatSlider(
    min=0.1, max=1.0, value=0.7, step=0.1, description="Layer Opacity"
)

# Shader Selection
shader_dropdown = widgets.Dropdown(
    options=nv.mesh_shader_names(), value="Phong", description="Shader"
)

# 4. Callbacks


def on_xray_change(change):
    """Handle xray change."""
    nv.opts.mesh_xray = change["new"]


def on_vis_change(change):
    """Handle vis change."""
    mesh_hemisphere.visible = change["new"]


def on_clear_layers_click(b):
    """Handle clear layers button."""
    mesh_hemisphere.layers = []


def on_add_layer_click(b):
    """Handle add layer button."""
    new_layer = MeshLayer(
        path=DATA_FOLDER / "BrainMesh_ICBM152.lh.motor.mz3",
        cal_min=0.5,
        cal_max=5.5,
        use_negative_cmap=True,
        opacity=0.7,
        colormap="warm",
        colormap_negative="winter",
    )
    # Update the list trait to trigger synchronization (instead of using .append)
    mesh_hemisphere.layers = [*mesh_hemisphere.layers, new_layer]
    print(mesh_hemisphere.layers)


def on_layer_thresh_change(change):
    """Handle layer thresh change."""
    if len(mesh_hemisphere.layers) > 0:
        # Update the first layer's cal_min
        mesh_hemisphere.layers[0].cal_min = change["new"]


def on_layer_alpha_change(change):
    """Handle layer alpha change."""
    if len(mesh_hemisphere.layers) > 0:
        # Update the first layer's opacity
        mesh_hemisphere.layers[0].opacity = change["new"]


def on_shader_change(change):
    """Handle shader change."""
    # Set shader for the first mesh
    nv.set_mesh_shader(mesh_hemisphere.id, change["new"])


# 5. Bind Callbacks
xray_slider.observe(on_xray_change, names="value")
vis_checkbox.observe(on_vis_change, names="value")
clear_layers_btn.on_click(on_clear_layers_click)
add_layer_btn.on_click(on_add_layer_click)
layer_thresh_slider.observe(on_layer_thresh_change, names="value")
layer_alpha_slider.observe(on_layer_alpha_change, names="value")
shader_dropdown.observe(on_shader_change, names="value")

# 6. Display All
controls_row1 = widgets.HBox([xray_slider, vis_checkbox, add_layer_btn])
controls_row2 = widgets.HBox(
    [layer_thresh_slider, layer_alpha_slider, clear_layers_btn]
)
controls_row3 = widgets.HBox([shader_dropdown])

layout = widgets.VBox([controls_row1, controls_row2, controls_row3, nv])

display(layout)